import deepxde as dde
import numpy as np
from matplotlib import pyplot as plt

"""----------------- Fluid parameters  --------------"""
rho = 1.2 #(kg/m^3)
mu = 1.8e-5 #(Pa·s)
nu = mu/rho #(m^2/s)

level = 4

# data
data_points = np.load('/home/aistudio/work/room/level' + str(level) + '/1_Fluent_results/level' + str(level) + '_fluent_results.npy') 
data = dde.data.DataSet(X_train=data_points[:, 0:3],
                    y_train=data_points[:, 3:7],
                    X_test=data_points[:, 0:3],
                    y_test=data_points[:, 3:7],
                    standardize=True,
                    ) 
Mean = data_points[:, 0:3].mean(axis=0)
SD = data_points[:, 0:3].std(axis=0)

"""------------------------------ Navier Stokes Equation(Time-dependent PDE) -----------------------"""
def standardize_NS_Equations_3D(x, y):
    """
    System of PDEs to be minimized: incompressible 2D Navier-Stokes equation 
    in the velocity-pressure formulations.

    """
    u = y[:, 0:1]
    v = y[:, 1:2]
    w = y[:, 2:3]
    # p = y[:, 3:4]
    du_x = dde.grad.jacobian(y, x, i=0, j=0) / SD[0]
    du_y = dde.grad.jacobian(y, x, i=0, j=1) / SD[1]
    du_z = dde.grad.jacobian(y, x, i=0, j=2) / SD[2]
    dv_x = dde.grad.jacobian(y, x, i=1, j=0) / SD[0]
    dv_y = dde.grad.jacobian(y, x, i=1, j=1) / SD[1]
    dv_z = dde.grad.jacobian(y, x, i=1, j=2) / SD[2]
    dw_x = dde.grad.jacobian(y, x, i=2, j=0) / SD[0]
    dw_y = dde.grad.jacobian(y, x, i=2, j=1) / SD[1]
    dw_z = dde.grad.jacobian(y, x, i=2, j=2) / SD[2]
    dp_x = dde.grad.jacobian(y, x, i=3, j=0) / SD[0]
    dp_y = dde.grad.jacobian(y, x, i=3, j=1) / SD[1]
    dp_z = dde.grad.jacobian(y, x, i=3, j=2) / SD[2]
    du_xx = dde.grad.hessian(y, x, component=0, i=0, j=0) / (SD[0]**2)
    du_yy = dde.grad.hessian(y, x, component=0, i=1, j=1) / (SD[1]**2)
    du_zz = dde.grad.hessian(y, x, component=0, i=2, j=2) / (SD[2]**2)
    dv_xx = dde.grad.hessian(y, x, component=1, i=0, j=0) / (SD[0]**2)
    dv_yy = dde.grad.hessian(y, x, component=1, i=1, j=1) / (SD[1]**2)
    dv_zz = dde.grad.hessian(y, x, component=1, i=2, j=2) / (SD[2]**2)
    dw_xx = dde.grad.hessian(y, x, component=2, i=0, j=0) / (SD[0]**2)
    dw_yy = dde.grad.hessian(y, x, component=2, i=1, j=1) / (SD[1]**2)
    dw_zz = dde.grad.hessian(y, x, component=2, i=2, j=2) / (SD[2]**2)
    continuity = du_x + dv_y + dw_z
    x_momentum = u * du_x + v * du_y + w * du_z + (1/rho) * dp_x - nu * (du_xx + du_yy + du_zz)
    y_momentum = u * dv_x + v * dv_y + w * dv_z + (1/rho) * dp_y - nu * (dv_xx + dv_yy + dv_zz)
    z_momentum = u * dw_x + v * dw_y + w * dw_z + (1/rho) * dp_z - nu * (dw_xx + dw_yy + dw_zz)
    return [continuity, x_momentum, y_momentum, z_momentum]

def restore_model(model_path):
    layer_size = [3] + [40] * 8 + [4]
    activation = "tanh"
    initializer = "Glorot uniform"
    net = dde.nn.FNN(layer_size, activation, initializer)
    model = dde.Model(data, net)
    model.compile("adam", lr=0.001)
    model.restore(save_path = model_path)
    return model
    
def plot_room_flow_field(model):
    import matplotlib as mpl
    from matplotlib import rcParams
    config = {"font.size": 6}
    rcParams.update(config)
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    
    '''------------------------------------------- Model validation ---------------------------------------------'''
    '''------------------Plotting 2------------------'''
    interface = dde.geometry.Rectangle([0, 0], [3, 3])
    interface_xy = interface.uniform_points(20000)

    test_data = np.hstack((interface_xy, 1.5 * np.ones_like(interface_xy[:, 0:1])))
    test_data2 = data.transform_inputs(test_data)

    # Model predictions generation
    u = model.predict(test_data2)[:, 0]
    v = model.predict(test_data2)[:, 1]
    w = model.predict(test_data2)[:, 2]
    p = model.predict(test_data2)[:, 3]
    velocity = (u**2 + v**2 + w**2)**0.5

    cmap_color = "jet"
    # plot
    fig, ax = plt.subplots(figsize=(3, 2), dpi=300)
    ax.tricontourf(test_data[:, 0], test_data[:, 1], velocity, levels=1000, cmap=cmap_color, vmin=0, vmax=1)
    plt.axis('scaled')
    norm =mpl.colors.Normalize(vmin=0, vmax=1)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Velocity Magnitude')
    # plt.savefig("./Figs/Velocity.png")
    plt.show()   
    
    fig, ax = plt.subplots(figsize=(3, 2), dpi=300)
    ax.tricontourf(test_data[:, 0], test_data[:, 1], p, levels=1000, cmap=cmap_color, vmin=10.2, vmax=11.0)
    plt.axis('scaled')
    norm =mpl.colors.Normalize(vmin=10.2, vmax=11.0)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[Pa]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('Pressure')
    # plt.savefig("./Figs/Pressure.png")
    plt.show()  
    
    fig, ax = plt.subplots(figsize=(3, 2), dpi=300)
    ax.tricontourf(test_data[:, 0], test_data[:, 1], u, levels=1000, cmap=cmap_color, vmin=-2.5, vmax=3.3)
    plt.axis('scaled')
    norm =mpl.colors.Normalize(vmin=-2.5, vmax=3.3)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('u')
    # plt.savefig("./Figs/u.png")
    plt.show()  

    fig, ax = plt.subplots(figsize=(3, 2), dpi=300)
    ax.tricontourf(test_data[:, 0], test_data[:, 1], v, levels=1000, cmap=cmap_color, vmin=-1, vmax=4)
    plt.axis('scaled')
    norm =mpl.colors.Normalize(vmin=-1, vmax=4)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('v')
    # plt.savefig("./Figs/v.png")
    plt.show()  

    fig, ax = plt.subplots(figsize=(3, 2), dpi=300)
    ax.tricontourf(test_data[:, 0], test_data[:, 1], w, levels=1000, cmap=cmap_color, vmin=-0.1, vmax=0.1)
    plt.axis('scaled')
    norm =mpl.colors.Normalize(vmin=-0.1, vmax=0.1)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('w')
    # plt.savefig("./Figs/w.png")
    plt.show() 
